import numpy as np
import scipy as sp

import common


class IceCover:
    """
    Model of ice cover in the Arctic water.

    Ice cover at a geographical point are described by the maximum speed of a 60 MW icebreaker.

    Attributes:
        lat: Latitude for each point in the spatial grid (deg).
        lon: Longitude for each point in the spatial grid (deg).
        speed_60mw: Speed of an 60 MW icebreaker for each point in the spatial grid (knots).
    """

    class Path:
        """
        Path between two points.

        A path is defined as a sequence of points from the spatial grid.

        Attributes:
            lat: Lattitude for each point in the path (deg).
            lon: Longitude for each point in the path (deg).
            speed_60mw: Speed of an 60 MW icebreaker for each point in the path (knots).
        """
        lat: list[float]
        lon: list[float]
        speed_60mw: list[float]
        part: list[float]
        length: float

        def __init__(self, lat: list[float], lon: list[float], speed_60mw: list[float], part: list[float],
                     length: float):
            self.lat = lat
            self.lon = lon
            self.speed_60mw = speed_60mw
            self.part = part
            self.length = length

        def __len__(self):
            return len(self.lat)

    lat: np.ndarray
    lon: np.ndarray
    speed_60mw: np.ndarray
    interpolator: sp.interpolate.NearestNDInterpolator

    def __init__(self, lat: np.ndarray, lon: np.ndarray, speed_60mw: np.ndarray):
        """
        Args:
            lat: Latitude for each point of the spatial grid (deg).
            lon: Longitude for each point of the spatial grid (deg).
            speed_60mw: Speed of an 60 MW icebreaker for each point in the spatial grid (knots).
        """

        # Save the data.
        self.lat = lat
        self.lon = lon
        self.speed_60mw = speed_60mw

        # Decompose the grid into a list of points.
        m, n = self.lat.shape
        X = np.block([self.lat.reshape((m * n, 1)), self.lon.reshape((m * n, 1))])
        Y = self.speed_60mw.reshape(m * n)

        # Compute the weight matrix of the graph. The weight of an edge connecting two points is
        # defined as the number of hours needed to pass the distance between the points.
        row = []
        col = []
        data = []
        for i0 in range(1, m - 1):
            for j0 in range(1, n - 1):
                k0 = np.ravel_multi_index((i0, j0), (m, n)).astype(int)
                s0 = speed_60mw[i0, j0]
                if s0 <= 0.0:
                    continue
                for i1, j1 in [(i0, j0 + 1), (i0 - 1, j0), (i0, j0 - 1), (i0 + 1, j0)]:
                    k1 = np.ravel_multi_index((i1, j1), (m, n)).astype(int)
                    s1 = speed_60mw[i1, j1]
                    if s1 <= 0.0:
                        continue
                    s01 = 0.5 * (s0 + s1)
                    length = common.path_lenghth(self.lat[i0, j0], self.lon[i0, j0],
                                                 self.lat[i1, j1], self.lon[i1, j1])
                    row.append(k0)
                    col.append(k1)
                    data.append(length / s01)
        W = sp.sparse.csr_array((data, (row, col)))

        # Find the shortest path for each pair of nodes.
        self.distances, self.predecessors =\
            sp.sparse.csgraph.shortest_path(W, method="auto", directed=False, return_predecessors=True)

        # Find the subset of points accessible from other points.
        self.mask = []
        for k in range(self.distances.shape[0] - 1):
            if np.isinf(self.distances[k + 1:, k]).all():
                continue
            self.mask.append(k)

        # Create a kd-tree for nearest neighbor lookup.
        self.kd_tree = sp.spatial.KDTree(X[self.mask, :])

        # Create a nearest-neighbors interpolator.
        self.interpolator = sp.interpolate.NearestNDInterpolator(X, Y)

    def update(self, speed_60mw: np.ndarray):
        """
        Update the speed values over the spatial grid.

        Args:
            speed_60mw: Speed of an 60 MW icebreaker for each point in the spatial grid (knots).
        """

        # Save the data.
        self.speed_60mw = speed_60mw

        # Decompose the grid into a list of points.
        X = np.block([self.lat.flat[:, np.newaxis], self.lon.flat[:, np.newaxis]])
        Y = self.speed_60mw.flatten()

        # Create a new interpolator.
        self.interpolator = sp.interpolate.NearestNDInterpolator(X, Y)

    def compute_cover(self, lat: float, lon: float) -> float:
        """
        Compute the ice cover at a geographical point.

        Args:
            lat: Latitude (deg).
            lon: Longitude (deg).

        Returns:
            The speed of a 60 MW icebreaker (knots).
        """
        return self.interpolator(lat, lon)

    def find_path(self, lat1: float, lon1: float, lat2: float, lon2: float) -> Path:
        """
        Find the fastest path between two points.

        Args:
            lat1: Latitude of the first point (deg).
            lon1: Longitude of the first point (deg).
            lat2: Latitude of the second point (deg).
            lon2: Longitude of the second point (deg).

        Returns:
            The path.
        """

        # Find the point in the grid nearest to the first point.
        d0, l0 = self.kd_tree.query((lat1, lon1), k=1)
        k0 = self.mask[l0]

        # Find the nearest accessible point in the grid.
        dn, ln = self.kd_tree.query((lat2, lon2), k=1)
        kn = self.mask[ln]

        # Is there any path between these points?
        if np.isinf(self.distances[k0, kn]):
            return IceCover.Path([], [], [], [], np.inf)

        # Reconstruct the path.
        points = [kn]
        k = kn
        while k != k0:
            k = self.predecessors[k0, k]
            points.append(k)
        points = list(reversed(points))
        n = len(points)
        lat = []
        lon = []
        for i in range(n):
            k = points[i]
            lat.append(self.lat.flat[k])
            lon.append(self.lon.flat[k])
        speed_mw60 = []
        for i in range(1, n):
            k1 = points[i - 1]
            k2 = points[i]
            s12 = 0.5 * (self.speed_60mw.flat[k1] + self.speed_60mw.flat[k2])
            speed_mw60.append(s12)
        length = 0
        part = [0.0]
        for i in range(1, n):
            lat1 = lat[i - 1]
            lon1 = lon[i - 1]
            lat2 = lat[i]
            lon2 = lon[i]
            length += common.path_lenghth(lat1, lon1, lat2, lon2)
            part.append(length)
        for i in range(1, n):
            part[i] /= length

        return IceCover.Path(lat, lon, speed_mw60, part, length)
